package telex;
/*
 * LEXComputer.java
 *
 * Copyright (c) 2007-2008 Colin Bayer, Douglas Downey, Oren Etzioni,
 *   University of Washington Computer Science and Engineering
 * 
 * See the file "COPYING" in the root directory of this distribution for
 * full licensing information.
 */

import java.util.ArrayList;
import java.util.List;
import java.util.Arrays;
import java.io.*;

public class LEXComputer {
	public static class RequiredQueryNotPresentException extends RuntimeException {
		private static final long serialVersionUID = -6757855544570834001L;
	}
	
	public LEXComputer(int aMaxN) {
		if (aMaxN <= 1) throw new IllegalArgumentException();
		
		mMaxN = aMaxN;
	}
	
	public boolean setDatabase(NGramDatabase aDB) {
		//if (aDB.getMaximumArity() < mMaxN) return false;
		
		mDB = aDB;
		return true;
	}
	
	public double computeDelta(String aWord) {
		// assumption: aWord is already capitalized.
		//String[] start = new String[]{"<S>", aWord};
		String[] start = new String[]{"#EOS#", aWord};
		String[] nonstart = new String[]{aWord};

                Query q_nonstart = mDB.executeQuery(new Query(nonstart));
		Query q_start = mDB.executeQuery(new Query(start));

		
		if (q_start.getResult() == -1 || q_nonstart.getResult() == -1) return -1;
		else return (double)q_start.getResult() / q_nonstart.getResult();
	}
	
	public double computeTau(double aNumExpt, String[] aA, String[] aB, String[] aC) {
		if (aA == null || aB == null || aC == null) throw new IllegalArgumentException();
		
		String[] merged = new String[aA.length + aB.length + aC.length];
		
		System.arraycopy(aA, 0, merged, 0, aA.length);
		System.arraycopy(aB, 0, merged, aA.length, aB.length);
		System.arraycopy(aC, 0, merged, aA.length + aB.length, aC.length);
		
		double num = composeProbability(merged);
		double denom_a = composeProbability(aA), denom_b = composeProbability(aB), 
			denom_c = composeProbability(aC);
		
		if (num > 0 || denom_a > 0 || denom_b > 0 || denom_c > 0) {
			System.out.println("Some query missing: " + num + " " + denom_a + " " + denom_b + " " + denom_c);
			return -1;
		}
		else if (num == 0) return 0;
		else {
			System.out.println("Tau for " + Arrays.toString(aA) + " and " + Arrays.toString(aC) + " across " + 
				Arrays.toString(aB) + " is exp(" + num + " * " + aNumExpt + " - (" + denom_a + " + " + denom_c + " + " + denom_b + "))");
			return Math.exp(num * aNumExpt - (denom_a + denom_b + denom_c));
		}
	}
	
	public static Sentence markCapsEntities(Sentence aSent) {
		for (int i = 0; i < aSent.length(); i++) {
			if (Character.isUpperCase(aSent.getWord(i).charAt(0))) {
				//aSent.setThresholds(i, Double.POSITIVE_INFINITY, Double.POSITIVE_INFINITY);
				aSent.setInNamedEntity(i);
			}
		}
		
		return aSent;
	}
	
	/* return log-prob of aS in the current database. */
	public double composeProbability(String[] aS) {
		double log_pr = 0;
		String deriv = "";
		
		if (aS.length <= mMaxN) {
			// probability can be computed from a single query.
			Query q = new Query(aS);
			
			if ((q = mDB.executeQuery(q)) == null) throw new RequiredQueryNotPresentException();
			
			if (q == null || q.getResult() < 0) {
				System.out.println("Internal error: query " + q + " not found in database");
				throw new RequiredQueryNotPresentException();
			} else if (q.getResult() == 0) {
				// if query did not occur in data, pr = 0.
				deriv = "0 results";
				log_pr = Double.NEGATIVE_INFINITY;
			} else {
				deriv = q.getResult() + " of " + mDB.getTotalCount(aS.length)[0] + " " + aS.length + "-grams";
				log_pr = Math.log(q.getResult()) - Math.log(mDB.getTotalCount(aS.length)[0]);
			}
		} else {
			// probability needs to be pieced together from multiple queries:
			// ex. p(abcde) ~= p(abcd) * p(e|bcd) = p(abcd) * p(bcde) / p(bcd)
			String[][] qs = decomposeString(aS, mMaxN).toArray(new String[0][]);
			double num = 0, denom = 0;
			
			deriv = "derived under independence assumption; ";
			
			for (String[] q : qs) {
				Query qobj = mDB.executeQuery(new Query(q));
				
				if (qobj == null || qobj.getResult() < 0) {
					System.out.println("Internal error: query " + q + " not found in database");
					throw new RequiredQueryNotPresentException();
				}
				
				// HACK: we should have a better way of determining numerator/denominator queries.
				if (q.length == mMaxN) {
					// numerator query.
					if (qobj.getResult() == 0) {
						num = 0;
						deriv = "numerator term " + qobj;
						break;
					} else {
						num += Math.log(qobj.getResult()) - Math.log(mDB.getTotalCount(q.length)[0]);
					}
					
					deriv += "numerator term " + qobj + ", ";
				} else {
					// denominator query.
					if (qobj.getResult() == 0) {
						denom = 0;
						deriv = "numerator term " + qobj;
						break;
					} else {
						denom += Math.log(qobj.getResult()) - Math.log(mDB.getTotalCount(q.length)[0]);
					}
					
					deriv += "denominator term " + qobj + ", ";
				}
			}
			
			if (num == 0 || denom == 0) return Double.NEGATIVE_INFINITY;
			log_pr = num - denom;
		}
		
		System.out.println("Log-prob of " + Arrays.toString(aS) + " is " + log_pr + " (" + deriv + ")");
		return log_pr;
	}
	
	private static List<String[]> decomposeString(String[] aS, int aMaxN) {
		ArrayList<String[]> rv = new ArrayList<String[]>();
		
		if (aS.length <= aMaxN) {
			rv.add(aS);
		} else {
			for (int i = 0; i < aS.length - aMaxN + 1; i++) {
				String[] n = new String[aMaxN];
				
				System.arraycopy(aS, i, n, 0, aMaxN);
				rv.add(n);
				
				//System.out.println("Added numerator " + Arrays.toString(n) + " for query " + Arrays.toString(aS));
				
				if (i != 0) {
					String[] d = new String[aMaxN - 1];
					
					System.arraycopy(aS, i, d, 0, aMaxN - 1);
					rv.add(d);
					
					//System.out.println("Added denominator " + Arrays.toString(d) + " for query " + Arrays.toString(aS));
				}
			}
		}
		
		return rv;
	}
	
	/* 
	 * requiredQueries() determines the queries required to calculate a merge operation between
	 * the entities aA and aC, separated by phrase aB.
	 */
	public Query[] requiredQueries(String[] aA, String[] aB, String[] aC) {
		int a_len = aA.length, b_len = aB.length, c_len = aC.length;
		ArrayList<String[]> s = new ArrayList<String[]>();
		
		String[] joined = new String[a_len + b_len + c_len];
		System.arraycopy(aA, 0, joined, 0, a_len);
		System.arraycopy(aB, 0, joined, a_len, b_len);
		System.arraycopy(aC, 0, joined, a_len + b_len, c_len);
		
		s.addAll(decomposeString(aA, mMaxN));
		s.addAll(decomposeString(aB, mMaxN));
		s.addAll(decomposeString(aC, mMaxN));
		s.addAll(decomposeString(joined, mMaxN));
		
		Query[] qs = new Query[s.size()];
		
		// FIXME: there are going to be a /lot/ of duplicated queries in the set, especially for
		// large entities.  we should remove duplicated queries from the array.
		// FIXME UPDATE: we handle this in QuerySorter.  This is no longer a priority.
		
		for (int i = 0; i < s.size(); i++) {
			qs[i] = new Query(s.get(i));
		}
		
		return qs;
	}
	
	private int mMaxN;
	private NGramDatabase mDB;
	
	/* testing code. */
	public static void main(String[] aArgs) {
		LEXComputer lx = new LEXComputer(5);
		try {
			IQFNGramDatabase db = new IQFNGramDatabase("D:\\Google ngrams\\data"); 
			
			db.addIQF(new IQFReader("O:\\unix\\homes\\iws\\vogon\\telex2\\in.file.iqf.sorted.resolved"));
			
			lx.setDatabase(db);
		} catch (Exception e) {
			System.out.println("Couldn't open IQFNGramDatabase: " + e);
			System.exit(1);
		}
		
		String[] p = new String[]{"Procter"};
		String[] n = new String[]{"and"};
		String[] g = new String[]{"Gamble"};
		
		System.out.println(Arrays.toString(lx.requiredQueries(p, n, g)));
		
		String[] hb = new String[]{"Honest","Bob"};
		String[] nt = new String[]{"and","the"};
		String[] ftdi = new String[]{"Factory","to","Dealer","Incentives"};
		
		System.out.println(Arrays.toString(lx.requiredQueries(hb, nt, ftdi)));
		
		System.out.println("Log-prob of 'and' is " + lx.composeProbability(n));
		
		BufferedReader stdin = new BufferedReader(new InputStreamReader(System.in));
		
		while (true) {
			System.out.print("> ");
			System.out.flush();
			
			String[] w;
			
			try {
				w = stdin.readLine().split(" ");
			} catch (Exception e) {
				System.out.println("readLine failed: " + e);
				break;
			}
			
			if (w[0].equalsIgnoreCase("prob")) {
				String[] args = new String[w.length - 1];
				
				System.arraycopy(w, 1, args, 0, w.length - 1);
				
				double prob = lx.composeProbability(args);
				
				System.out.println("Log-prob of " + Arrays.toString(args) + " is " + prob);
			} else if (w[0].equalsIgnoreCase("delta")) {
				double delta = lx.computeDelta(w[1]);
				
				System.out.println("Delta threshold of '" + w[1] + "' is " + delta);
			} else if (w[0].equalsIgnoreCase("pmi") || w[0].equalsIgnoreCase("scp")) {
				int expt = 0;
				
				if (w[0].equalsIgnoreCase("pmi")) expt = 1;
				else expt = 3;
				
				// identify breaks between three argument sets.
				String[][] args = new String[3][];
				int i = 0;
				
				for (int j = 1; j < w.length; j++) {
					// this is slow as hell.
					if (w[j].equals("||")) { i++; continue; }
					
					if (args[i] == null) args[i] = new String[]{w[j]};
					else {
						String[] new_args = new String[args[i].length + 1];
						
						System.arraycopy(args[i], 0, new_args, 0, args[i].length);
						new_args[args[i].length] = w[j];
						
						args[i] = new_args;
					}
				}
				
				if (i != 2) System.out.println("Too few arguments.");
				double tau = lx.computeTau(expt, args[0], args[1], args[2]);
				
				System.out.println(w[0].toUpperCase() + " of " + Arrays.toString(args[0]) + " and " +
					Arrays.toString(args[2]) + " across " + Arrays.toString(args[1]) + " is " + tau);
			} else if (w[0].equalsIgnoreCase("quit")) break;
			else System.out.println("Invalid command '" + w[0] + "'.");
		}

	}
}
